import random
from tvm import relax
from tvm.relax.testing import nn
from tvm.script import ir as I, relax as R
import tvm

class RelaxFuseOpGen:
    def __init__(self):
        self._candidate_map = {
            'S': ['U', 'V', 'e'],
            'U': ['iU', 'ir', 'S'],  
            'V': ['cU']
        }

    def gen(self, n_round=10):
        ret = 'S'
        i = 0
        while ret[-1].isupper():
            candidates = self._candidate_map[ret[-1]].copy()
            if i < n_round and 'e' in candidates:
                candidates.remove('e')
            elif i >= n_round and 'e' in candidates:
                candidates = ['e']
            out = random.choice(candidates)
            ret = ret[:-1] + out
            i += 1
        return ret if ret[-1] != 'e' else ret[:-1]

    def gen_ir(self, n_round=10):
        expr_code = self.gen(n_round)
        print(f"Generated Expr Sequence: {expr_code}")

        bb = relax.BlockBuilder()
        with bb.function("main"):
            x = relax.Var("x", [1, 3, 128, 128], dtype="float32")
            w = relax.Var("w", [8, 3, 5, 5], dtype="float32")
            gv = x
            for v in expr_code:
                if v == 'i':
                    gv = random.choice(self.injective_ops())(bb, gv)
                elif v == 'r':
                    gv = random.choice(self.reduce_ops())(bb, gv)
                elif v == 'c':
                    gv = random.choice(self.out_ewise_fusable_ops())(bb, gv, w)
            bb.emit_func_output(gv)
        return tvm.IRModule.from_expr(bb.get())

    def injective_ops(self):
        return [
            lambda bb, x: bb.emit(relax.op.sqrt(x)),
            lambda bb, x: bb.emit(relax.op.relu(x)),
            lambda bb, x: bb.emit(relax.op.abs(x)),
            lambda bb, x: bb.emit(relax.op.ceil(x)),
            lambda bb, x: bb.emit(relax.op.floor(x)),
            lambda bb, x: bb.emit(relax.op.exp(x)),
            lambda bb, x: bb.emit(relax.op.pad(x, pad_width=[[0,0],[0,0],[1,1],[2,2]])),
        ]

    def reduce_ops(self):
        return [
            lambda bb, x: bb.emit(relax.op.sum(x)),
            lambda bb, x: bb.emit(relax.op.max(x)),
            lambda bb, x: bb.emit(relax.op.min(x)),
        ]

    def out_ewise_fusable_ops(self):
        return [
            lambda bb, x, w: bb.emit(relax.op.nn.conv2d(x, w, kernel_size=(5,5), padding=(1,1), out_dtype="float32")),
            lambda bb, x, w: bb.emit(relax.op.nn.conv2d_transpose(x, w, kernel_size=(5,5), padding=(1,1), out_dtype="float32")),
            lambda bb, x, _: bb.emit(relax.op.nn.max_pool2d(x, pool_size=(3,3))),
            lambda bb, x, _: bb.emit(relax.op.nn.avg_pool2d(x, pool_size=(3,3))),
        ]
